Make Vertex AI evaluation optional to avoid litellm bloat by Vamsi-klu · Pull Request #69749 · apache/airflow

Vamsi-klu · 2026-07-11T06:54:04Z

Problem

apache-airflow-providers-google pulled google-cloud-aiplatform[evaluation] into the base provider install. That extra brings in evaluation-only dependencies such as litellm, scikit-learn, huggingface-hub, and tokenizers for users who may only need unrelated Google services.

The only code path that needs vertexai.preview.evaluation is Vertex AI model evaluation through GenerativeModelHook.get_eval_task(), GenerativeModelHook.run_evaluation(), and RunEvaluationOperator.

What Changed

Changed the base dependency to google-cloud-aiplatform>=1.155.0.
Added an evaluation provider extra that installs google-cloud-aiplatform[evaluation]>=1.155.0.
Added an import guard that raises AirflowOptionalProviderFeatureException with an install instruction when evaluation APIs are used without the extra.
Added a provider changelog note for the breaking dependency split.
Updated generated provider dependency docs and uv.lock.
Added a follow-up regression test for the missing-extra path and documented the evaluation extra in the Vertex AI operator docs.

Why The Follow-Up Commit Exists

The feedback was correct: the original PR proved the opt-in path but did not directly test the failure mode when the extra is missing, did not document the extra next to RunEvaluationOperator, and left the lockfile inconsistent with the dependency split.

Impact

Base Google provider installs no longer pull the evaluation dependency stack.
Users of RunEvaluationOperator or direct Vertex AI evaluation hook methods must install apache-airflow-providers-google[evaluation].
Users who do not use evaluation APIs should see no runtime behavior change.

Fixes: #69323

Testing

UV_CACHE_DIR=/tmp/uv-cache-69749 uv run --python 3.11 ruff format providers/google/tests/unit/google/cloud/hooks/vertex_ai/test_generative_model_optional_evaluation.py
UV_CACHE_DIR=/tmp/uv-cache-69749 uv run --python 3.11 ruff check --fix providers/google/tests/unit/google/cloud/hooks/vertex_ai/test_generative_model_optional_evaluation.py
AIRFLOW_HOME=/tmp/airflow-home-69749 UV_CACHE_DIR=/tmp/uv-cache-69749 uv run --python 3.11 --project providers/google pytest providers/google/tests/unit/google/cloud/hooks/vertex_ai/test_generative_model_optional_evaluation.py --with-db-init -xvs
Result: 2 passed, 1 warning
AIRFLOW_HOME=/tmp/airflow-home-69749 UV_CACHE_DIR=/tmp/uv-cache-69749 uv run --python 3.11 --project providers/google --extra evaluation pytest providers/google/tests/unit/google/cloud/hooks/vertex_ai/test_generative_model.py providers/google/tests/unit/google/cloud/operators/vertex_ai/test_generative_model.py -xvs
Result: 2 passed, 1 warning
Commit hooks passed for the follow-up commit.

Was generative AI tooling used to co-author this PR?

Yes — Codex (GPT-5)

Generated-by: Codex (GPT-5) following the guidelines

Drafted-by: Codex (GPT-5); reviewed by @Vamsi-klu before posting

Vamsi-klu · 2026-07-11T06:55:05Z

Reviewers: @shahar1 (google provider CODEOWNER). Breaking change: evaluation extra now optional to avoid litellm bloat. Changelog entry added, provider docs regenerated via update-providers-dependencies.

Fixes #69323

Drafted-by: Muse Spark 1.1; reviewed by @Vamsi-klu before posting

github-actions · 2026-07-13T22:39:44Z

uv.lock on main just moved via #71039 ("Record the relaxed types-paramiko bound in uv.lock"), commit 381339d and this PR currently conflicts.

Quickest fix:

git fetch upstream main && git rebase upstream/main
rm uv.lock && uv lock
git add uv.lock && git rebase --continue
git push --force-with-lease

Automated nudge — ignore if you're not ready to rebase. This comment is updated in place on future uv.lock bumps.

potiuk · 2026-07-20T11:54:44Z

@Vamsi-klu This PR has been converted to draft because it does not yet meet our Pull Request quality criteria.

Issues found:

❌ Merge conflicts: This PR has merge conflicts with main and cannot be merged as-is. Rebase locally and resolve them: git fetch upstream main && git rebase upstream/main, resolve the conflicts, then force-push.

Note: Your branch is 147 commits behind main. Some check failures may be caused by changes in the base branch rather than by your PR. Please rebase your branch and push again to get up-to-date CI results.

What to do next:

Fix each issue listed above.
Make sure static checks pass locally (prek run --from-ref main --stage pre-commit).
Mark the PR as "Ready for review" when you're done.

Converting a PR to draft is not a rejection — it is an invitation to bring the PR up to the project's standards so that maintainer review time is spent productively. There is no rush — take your time and work at your own pace. We appreciate your contribution and are happy to wait for updates. If you have questions, feel free to ask on the Airflow Slack.

Note: This comment was drafted by an AI-assisted triage tool and may contain mistakes. Once you have addressed the points above, an Apache Airflow maintainer — a real person — will take the next look at your PR. We use this two-stage triage process so that our maintainers' limited time is spent where it matters most: the conversation with you.

Base dependency on google-cloud-aiplatform[evaluation] forced litellm, scikit-learn, tokenizers for all users. Move evaluation extra behind optional provider extra 'evaluation'. Add lazy import guard with AirflowOptionalProviderFeatureException in GenerativeModelHook. Users needing RunEvaluationOperator should install apache-airflow-providers-google[evaluation]. Fixes: apache#69323

Use TYPE_CHECKING guard and Any fallback to satisfy mypy when evaluation extra not installed. Runtime guard still raises AirflowOptionalProviderFeatureException before using None types. Fixes mypy failure in CI for apache#69323

The provider now keeps evaluation dependencies out of the base install, so the missing-extra path needs regression coverage and operator docs that show users how to opt back in.

Reloading the hook module rebound GenerativeModelHook, so tests that imported the class at collection time patched a different class object and hit the real EvalTask. Patching the import-guard variable tests the same behavior without touching module state.

aaron-y-chen

Thanks for the PR!

aaron-y-chen · 2026-08-06T02:07:21Z


 # For no Pydantic environment, we need to skip the tests
 pytest.importorskip("google.cloud.aiplatform_v1")
+pytest.importorskip("vertexai.preview.evaluation")


I noticed that the CI image installs via uv sync --all-packages --group ci-image (Dockerfile.ci, no --all-extras), so this skip fires there, and the only tests in these two files are test_run_evaluation and TestVertexAIRunEvaluationOperator::test_execute, i.e. exactly the paths this PR guards. Is that intended?

Good catch, thanks for pointing this out! You are right, this was not intended.

The pytest.importorskip("vertexai.preview.evaluation") in those two files would fire in the default CI image since it uses uv sync --all-packages --group ci-image without extras, and those files only contain test_run_evaluation and test_execute. So we would end up with no coverage for the evaluation path in CI.

I pushed a fix in 4c25fa3:

Removed the importorskip for vertexai.preview.evaluation from both test_generative_model.py files

Wrapped the MetricPromptTemplateExamples import in a try/except and fall back to a MagicMock with the Pointwise values when the evaluation extra is not installed. The tests already mock get_eval_task and get_generative_model, so this lets the wiring be tested without needing the real litellm stack

Also cleaned up the leftover merge markers in changelog.rst

Verified locally that pytest --collect-only now shows 2 tests collected and both pass, and the test_generative_model_optional_evaluation.py tests for the missing extra still pass. Let me know if you would prefer to keep the skip and instead run these tests in a job with the evaluation extra, happy to adjust.

Address review comment that pytest.importorskip('vertexai.preview.evaluation') caused the only tests in two files to be skipped in the default CI image which installs via uv sync --all-packages --group ci-image without extras. Replace the skip with a try/except import that falls back to a MagicMock providing MetricPromptTemplateExamples.Pointwise attributes. Tests now execute with mocks in default CI, while the missing-extra path remains covered by test_generative_model_optional_evaluation. Also fix changelog merge markers left from previous rebase.

boring-cyborg Bot added area:providers kind:documentation provider:google Google (including GCP) related issues labels Jul 11, 2026

Vamsi-klu force-pushed the fix/google-evaluation-optional-69323 branch from 8d5aab9 to b54b2d6 Compare July 12, 2026 19:56

Vamsi-klu marked this pull request as ready for review July 13, 2026 03:36

Vamsi-klu requested a review from shahar1 as a code owner July 13, 2026 03:36

potiuk marked this pull request as draft July 20, 2026 11:54

Vamsi-klu marked this pull request as ready for review July 20, 2026 12:03

Vamsi-klu force-pushed the fix/google-evaluation-optional-69323 branch from 1506a46 to ce7370f Compare July 20, 2026 14:07

Vamsi-klu force-pushed the fix/google-evaluation-optional-69323 branch 2 times, most recently from 5820c69 to 9ee89ab Compare August 1, 2026 04:36

Vamsi-klu and others added 6 commits August 4, 2026 06:37

Fix mypy for optional evaluation import

cd4c659

Use TYPE_CHECKING guard and Any fallback to satisfy mypy when evaluation extra not installed. Runtime guard still raises AirflowOptionalProviderFeatureException before using None types. Fixes mypy failure in CI for apache#69323

Document optional Vertex AI evaluation dependency

a913bdc

The provider now keeps evaluation dependencies out of the base install, so the missing-extra path needs regression coverage and operator docs that show users how to opt back in.

Fix CI failures after optional Vertex AI evaluation change

ac4c0bf

Fix optional evaluation test typing

32e7212

Vamsi-klu force-pushed the fix/google-evaluation-optional-69323 branch from 9ee89ab to 1141e4f Compare August 4, 2026 06:55

aaron-y-chen reviewed Aug 6, 2026

View reviewed changes

Navigation Menu

Uh oh!

Make Vertex AI evaluation optional to avoid litellm bloat - #69749

Make Vertex AI evaluation optional to avoid litellm bloat#69749
Vamsi-klu wants to merge 7 commits into
apache:mainfrom
Vamsi-klu:fix/google-evaluation-optional-69323

Vamsi-klu commented Jul 11, 2026 •

edited

Loading

Uh oh!

Vamsi-klu commented Jul 11, 2026

Uh oh!

github-actions Bot commented Jul 13, 2026 •

edited

Loading

Uh oh!

potiuk commented Jul 20, 2026

Uh oh!

aaron-y-chen left a comment

Uh oh!

aaron-y-chen Aug 6, 2026

Uh oh!

Vamsi-klu Aug 6, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Uh oh!

Conversation

Vamsi-klu commented Jul 11, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Problem

What Changed

Why The Follow-Up Commit Exists

Impact

Testing

Was generative AI tooling used to co-author this PR?

Uh oh!

Vamsi-klu commented Jul 11, 2026

Uh oh!

github-actions Bot commented Jul 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

potiuk commented Jul 20, 2026

Uh oh!

aaron-y-chen left a comment

Choose a reason for hiding this comment

Uh oh!

aaron-y-chen Aug 6, 2026

Choose a reason for hiding this comment

Uh oh!

Vamsi-klu Aug 6, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Vamsi-klu commented Jul 11, 2026 •

edited

Loading

github-actions Bot commented Jul 13, 2026 •

edited

Loading